# -*- coding: utf-8 -*-
import numpy as np
# import scipy.signal as signal
from scipy import signal
import pywt
import math
import copy
#matplotlib.rcParams['font.sans-serif'] = ['SimHei'] # 指定默认字体
#matplotlib.rcParams['axes.unicode_minus'] = False   # 正确显示符号

def load_txt(txtpath):
    # 读取两列的脑电信号数据
    data = np.loadtxt(fname=txtpath,dtype=float)
    data = data.reshape((np.shape(data)[0],-1))
    print('data',data.shape)
    return data

def remove_drift(data):
    data = signal.detrend(data=data,axis=0,type='linear')
    # axis=0按列去漂移
    return data

def band_pass(data):
    fs = 512
    f0 = 256
    f1 = 0.1
    f2 = 50
    # 对于模拟滤波器，wp和ws单位是rad/s
    # 圆频率（rad/s) = 2π*频率（Hz）
    # 我们的输入是离散的点，所以要设计数字滤波器

    # [b,a]= signal.butter(N=6,Wn=[wp11,wp22],btype='bandpass',analog=True)
    # b, a = signal.bilinear(b, a, fs=fs)
    # w2, h2 = signal.freqz(b, a, worN=1000)

    # [b, a] = signal.butter(N=6, Wn=[(fs*np.arctan(math.pi*f1/fs)/math.pi)/f0,(fs*np.arctan(math.pi*f2/fs)/math.pi)/f0])
    # [b, a] = signal.butter(N=6, Wn=[f1/f0,f2/f0],btype='bandpass')

    # [b, a] = signal.iirdesign(wp=[0.1/f0, 50/f0], ws=[0.05/f0, 60/f0], gpass=2, gstop=40,ftype='butter',output='ba')

    [b, a] = signal.iirfilter(N=6, Wn=[50/f0], rp=3, rs=60,btype='lowpass', ftype='butter',output='ba',analog=False)

    # w, h = signal.freqz(b, a)
    # power = 20 * np.log10(np.clip(np.abs(h), 1e-8, 1e100))
    # pl.subplot(111)
    # pl.plot(w / np.pi * 512 / 2, power)
    # pl.title(u"made by freqz ")
    # # pl.ylim(-100, 20)
    # pl.ylabel(u"gain(dB)")
    # pl.xlabel(u"frequency(Hz)")
    # pl.show()

    data = signal.lfilter(b=b,a=a,x=data,axis=0)
    return data

def limit_value(data):
    high_limit = 200
    low_limit = -200
    row,col = data.shape
    for i in range(row):
        for j in range(col):
            if data[i][j]>high_limit:
                data[i][j] = high_limit
            elif data[i][j]<low_limit:
                data[i][j] = low_limit
    return data

def remove_eyedata_single(data):
    fs = 512
    f0 = 256

    Left_Latency = math.floor(0.05 * fs)
    # print('Left_Latency',Left_Latency)
    # 25

    Right_Latency = math.floor(0.3 * fs)
    # print('Right_Latency', Right_Latency)
    # 153

    stand_min = 100

    eye_time = 10               # 提取眼电的窗口时间,可修改
    PointNumber = fs * eye_time #  提取眼电的窗口所需数据点数
    # print('PointNumber',PointNumber)
    # 5120

    blink_number = int(PointNumber / 64)
    # print('blink_number',blink_number)
    # 10秒内总共眨眼次数，每秒眨眼不能超过8次。
    # 80

    data_number = Left_Latency + Right_Latency+1
    # print('data_number',data_number)
    # 179

    eye_data = np.zeros((blink_number, data_number))
    hd_Data = data[0:PointNumber]
    #取前10s中的5120个数据点

    [b,a] = signal.iirfilter(N=6,Wn=[10/f0],rp=3, rs=60,btype='lowpass',output='ba',ftype='butter')
    EYE_Data = signal.lfilter(b=b,a=a,x=hd_Data,axis=0)
    # EYE_Data是前10s中的5120个数据点滤波后的波形

    Row = 0
    point = Left_Latency

    while point <= PointNumber:
        if (point - Left_Latency>=0) and (point + Right_Latency <= PointNumber-1) and (EYE_Data[point] >= stand_min):
            eye_data[Row,:] = EYE_Data[(point - Left_Latency):(point + Right_Latency+1)]
            Row += 1
            point += Right_Latency + 1 + Left_Latency
            # 略过这179个点
        else:
            point += 1

    row,col=eye_data.shape
    index = []
    for i in range(row):
        if sum(abs(eye_data[i,:])) != 0 :
            index.append(i)
    eye_data = eye_data[index,:]

    Standard_EYE = np.mean(eye_data,0)
    # 获取标准眼电
    # print('Standard_EYE',Standard_EYE.shape,type(Standard_EYE))
    # (179,) ndarray

    # 绘制标准眼电
    #Len = len(Standard_EYE)
    #print('Len',Len)
    #time_len = int(1000 * (Len / fs))
    #print('time_len',time_len)

    #t1 = np.linspace(start=0, stop=time_len, num=Len)
    #print('t1',t1.shape)

    #plt.plot(t1, Standard_EYE)
    #plt.xlabel('t(ms)')
    #plt.ylabel('电位（uV）')
    #plt.title('标准眼电')
    #plt.show()

    [b, a] = signal.iirfilter(N=6, Wn=[30 / f0], rp=3, rs=60, btype='lowpass', ftype='butter', output='ba',analog=False)
    EEG_Data = signal.lfilter(b=b,a=a,x=data)
    EEG_Data = EEG_Data.reshape((-1,1))
    # EEG_Data (247384, 1) <class 'numpy.ndarray'>
    # print('EEG_Data',EEG_Data.shape,type(EEG_Data))
    # 对原始数据进行滤波

    EEG_Len = len(EEG_Data)
    # print('EEG_Len',EEG_Len)

    EEG_STP = EEG_Len - Right_Latency
    # print('EEG_STP', EEG_STP)
    # EEG_STP 247231

    EEG_Window = int(Left_Latency + Right_Latency + 1)
    # print('EEG_Window', EEG_Window)
    # EEG_Window 179

    NoEYE_Data = np.zeros((EEG_Len,1))

    item_Flag = Left_Latency
    item = item_Flag

    while item <= EEG_STP:
        if (EEG_Data[item] >= stand_min) and (item >= item_Flag):
            for i in range(EEG_Window):
                NoEYE_Data[item - Left_Latency + i - 1,:] = EEG_Data[item - Left_Latency + i - 1,:] - Standard_EYE[i]

            item += Right_Latency + 1
            item_Flag = item + Right_Latency + Left_Latency
        else:
            NoEYE_Data[item,:] = EEG_Data[item,:]
            item += 1


    # print('NoEYE_Data',NoEYE_Data.shape)
    return NoEYE_Data

def remove_eyedata(data):
    row,col = data.shape
    res = []
    for j in range(col):
        # data_temp = data[j,:]
        temp = remove_eyedata_single(data[:,j])
        if res == [] :
            res = temp
        else:
            res = np.hstack((res,temp))
    print('remove_eyedata',res.shape,type(res))
    return res

# ToDo
def xiaobo_quzao(data):
    # 小波变换分解，抛弃部分系数低的，重构信号
    res = pywt.wavedec(data=data,wavelet='db6',level=5,axis=0)
    print('res',len(res[0]))

    return data

def remove_mean(data):
    row,col = data.shape
    for i in range(col):
        mean = sum(data[:,i]) / row
        for j in range(row):
            data[j,i] -= mean
    return data

def data_pre_op_entry(data):
    data_remove_drift = remove_drift(data)
    data_bandpass = band_pass(data_remove_drift)

    data_limit = limit_value(data_bandpass)

    data_remove_eyedata = remove_eyedata(data_limit)

    data_wavelet = xiaobo_quzao(data_remove_eyedata)
    
    return data_wavelet;

def extract_Alpha(data):
    data_A = copy.copy(data)
    fs = 512
    f0 = 256
    [b, a] = signal.iirfilter(N=6, Wn=[8/f0,13/f0], rp=3, rs=60,btype='bandpass', ftype='butter',output='ba',analog=False)
    data = signal.lfilter(b=b, a=a, x=data_A, axis=0)
    return data_A

def extract_BeiTa(data):
    data_B = copy.copy(data)
    fs = 512
    f0 = 256
    [b, a] = signal.iirfilter(N=6, Wn=[6/f0,14/f0], rp=3, rs=60,btype='bandpass', ftype='butter',output='ba',analog=False)
    data = signal.lfilter(b=b, a=a, x=data_B, axis=0)
    return data_B

def extract_DeiTa(data):
    data_D = copy.copy(data)
    fs = 512
    f0 = 256
    [b, a] = signal.iirfilter(N=6, Wn=[0.5/f0,3/f0], rp=3, rs=60,btype='bandpass', ftype='butter',output='ba',analog=False)
    data = signal.lfilter(b=b, a=a, x=data_D, axis=0)
    return data_D

def extract_GaMa(data):
    data_G = copy.copy(data)
    fs = 512
    f0 = 256
    [b, a] = signal.iirfilter(N=6, Wn=[31/f0,80/f0], rp=3, rs=60,btype='bandpass', ftype='butter',output='ba',analog=False)
    data = signal.lfilter(b=b, a=a, x=data_G, axis=0)
    return data_G

def extract_SeiTa(data):
    data_S = copy.copy(data)
    fs = 512
    f0 = 256
    [b, a] = signal.iirfilter(N=6, Wn=[4/f0,7/f0], rp=3, rs=60,btype='bandpass', ftype='butter',output='ba',analog=False)
    data = signal.lfilter(b=b, a=a, x=data_S, axis=0)
    return data_S

def cal_power_(data):
    res = []
    row,col = data.shape
    for i in range(col):
        temp = data[:,i]
        temp_res=0
        for j in temp:
            temp_res += j*j
        res.append(temp_res)
    res = np.array(res)
    res.reshape((1,-1))
    return res

def test_print(data):
    return data;

if __name__=='__main__':
    data = load_txt('C:/Users/Administrator/Desktop/two_channel.txt')
    # data = load_txt('new_s2.txt')
    '''
    data_remove_drift = remove_drift(data)
    data_bandpass = band_pass(data_remove_drift)

    data_limit = limit_value(data_bandpass)

    data_remove_eyedata = remove_eyedata(data_limit)

    data_wavelet = xiaobo_quzao(data_remove_eyedata)
'''
    print(data_pre_op_entry(data))

    # row,col = data.shape
    # row = 5000
    # fig = plt.figure()
    # ax = fig.add_subplot(111)
    # x = list(range(row))
    # plt.plot(x,data_remove_eyedata[0:row,0])
    # plt.show()